For example, I have two lists of the same length:
a = [[1,2],[2,2],[3,3],[4,2],[5,6]]
b = [1,2,2,3,1]
I want to have a function such that
func(a,b,1) = [[1,2],[5,6]]
func(a,b,2) = [[2,2],[3,3]]
func(b,b,2) = [2,2]
What the function does is return a list of a's elements, whose corresponding elements of the same index in list b equal to the third argument.
In Matlab I will do something as easy as a(b==1), a(b==2), b(b==2). What is the most efficient way to achieve this in Python?
DSM :
If you want very Matlab-like functionality, you could use numpy:\n\n>>> import numpy\n>>> a = [[1,2],[2,2],[3,3],[4,2],[5,6]]\n>>> b = [1,2,2,3,1]\n>>> a = numpy.array(a)\n>>> b = numpy.array(b)\n>>> a[b==1]\narray([[1, 2],\n [5, 6]])\n>>> a[b==2]\narray([[2, 2],\n [3, 3]])\n>>> b[b==2]\narray([2, 2])\n\n\nFailing that, I'd probably simply use a list comprehension:\n\n>>> [i for i,j in zip(a,b) if j == 1]\n[[1, 2], [5, 6]]\n>>> [i for i,j in zip(a,b) if j == 2]\n[[2, 2], [3, 3]]\n>>> [i for i,j in zip(b,b) if j == 2]\n[2, 2]\n\n\nIt'd be trivial to wrap this in a function:\n\n>>> def func(a,b,x):\n... return [i for i,j in zip(a,b) if j == x]\n... \n>>> func(a,b,2)\n[[2, 2], [3, 3]]\n>>> \n",
2012-06-21T20:32:18
Gareth Latty :
I would use itertools.compress() along with a generator expression to do this:\n\ndef func(a, b, c):\n return itertools.compress(a, (x == c for x in b))\n\n\nNote that this will return a generator. If you need a list, wrap it in a list() call.\n\n>>> import itertools\n>>> def func(a, b, c):\n... return list(itertools.compress(a, (x == c for x in b)))\n... \n>>> a = [[1,2],[2,2],[3,3],[4,2],[5,6]]\n>>> b = [1,2,2,3,1]\n>>> func(a, b, 1)\n[[1, 2], [5, 6]]\n>>> func(a, b, 2)\n[[2, 2], [3, 3]]\n>>> func(b, b, 2)\n[2, 2]\n\n\nIt's also worth noting this should be nice and fast - itertools is designed to be a fast module.",
2012-06-21T20:35:21
Junuxx :
This first finds the indices where b equals the target values ([i for i, x in enumerate(b) if x==c]), then uses those indices to get the elements you want out of a:\n\ndef func(a,b,c):\n return [a[j] for j in [i for i, x in enumerate(b) if x==c]]\n\n>>> func(a,b,1)\n[[1, 2], [5, 6]]\n>>> func(a,b,2)\n[[2, 2], [3, 3]]\n>>> func(b,b,2)\n[2, 2]\n",
2012-06-21T20:31:14